import xml.etree.ElementTree as ET
import os
import json

coco = dict()
coco['images'] = []
coco['type'] = 'instances'
coco['annotations'] = []
coco['categories'] = []

category_set = dict()
image_set = set()

category_item_id = -1
image_id = 20180000000
annotation_id = 0


def addCatItem(name):
    global category_item_id
    category_item = dict()
    category_item['supercategory'] = 'none'
    category_item_id += 1
    category_item['id'] = category_item_id
    category_item['name'] = name
    coco['categories'].append(category_item)
    category_set[name] = category_item_id
    return category_item_id


def addImgItem(file_name, size):
    global image_id
    if file_name is None:
        raise Exception('Could not find filename tag in xml file.')
    if size['width'] is None:
        raise Exception('Could not find width tag in xml file.')
    if size['height'] is None:
        raise Exception('Could not find height tag in xml file.')
    image_id += 1
    image_item = dict()
    image_item['id'] = image_id
    image_item['file_name'] = file_name
    image_item['width'] = size['width']
    image_item['height'] = size['height']
    coco['images'].append(image_item)
    image_set.add(file_name)
    return image_id


def addAnnoItem(object_name, image_id, category_id, bbox):
    global annotation_id
    annotation_item = dict()
    annotation_item['segmentation'] = []
    seg = []
    # bbox[] is x,y,w,h
    # left_top
    seg.append(bbox[0])
    seg.append(bbox[1])
    # left_bottom
    seg.append(bbox[0])
    seg.append(bbox[1] + bbox[3])
    # right_bottom
    seg.append(bbox[0] + bbox[2])
    seg.append(bbox[1] + bbox[3])
    # right_top
    seg.append(bbox[0] + bbox[2])
    seg.append(bbox[1])

    annotation_item['segmentation'].append(seg)

    annotation_item['area'] = bbox[2] * bbox[3]
    annotation_item['iscrowd'] = 0
    annotation_item['ignore'] = 0
    annotation_item['image_id'] = image_id
    annotation_item['bbox'] = bbox
    annotation_item['category_id'] = category_id
    annotation_id += 1
    annotation_item['id'] = annotation_id
    coco['annotations'].append(annotation_item)


def _read_image_ids(image_sets_file):
    ids = []
    with open(image_sets_file) as f:
        for line in f:
            ids.append(line.rstrip())
    return ids


"""通过txt文件生成"""
# # split ='train' 'va' 'trainval' 'test'
# def parseXmlFiles_by_txt(data_dir, json_save_path, split='train'):
#     print("hello")
#     labelfile = split + ".txt"
#     image_sets_file = data_dir + "/ImageSets/Main/" + labelfile
#     ids = _read_image_ids(image_sets_file)
#
#     for _id in ids:
#         xml_file = data_dir + f"/Annotations/{_id}.xml"
#
#         bndbox = dict()
#         size = dict()
#         current_image_id = None
#         current_category_id = None
#         file_name = None
#         size['width'] = None
#         size['height'] = None
#         size['depth'] = None
#
#         tree = ET.parse(xml_file)
#         root = tree.getroot()
#         if root.tag != 'annotation':
#             raise Exception('pascal voc xml root element should be annotation, rather than {}'.format(root.tag))
#
#         # elem is <folder>, <filename>, <size>, <object>
#         for elem in root:
#             current_parent = elem.tag
#             current_sub = None
#             object_name = None
#
#             if elem.tag == 'folder':
#                 continue
#
#             if elem.tag == 'filename':
#                 file_name = elem.text
#                 if file_name in category_set:
#                     raise Exception('file_name duplicated')
#
#             # 如果获取了图片名和图片尺寸,则给予该图片对应id
#             if current_image_id is None and file_name is not None and size['width'] is not None:
#                 if file_name not in image_set: # 图片名不能与之前重复
#                     current_image_id = addImgItem(file_name, size)
#                     print('add image with {} and {}'.format(file_name, size))
#                 else:
#                     raise Exception('duplicated image: {}'.format(file_name))
#
#             # subelem is <width>, <height>, <depth>, <name>, <bndbox>
#             for subelem in elem: # 找子类
#                 bndbox['xmin'] = None
#                 bndbox['xmax'] = None
#                 bndbox['ymin'] = None
#                 bndbox['ymax'] = None
#
#                 current_sub = subelem.tag
#                 if current_parent == 'object' and subelem.tag == 'name':
#                     object_name = subelem.text
#                     if object_name not in category_set:
#                         current_category_id = addCatItem(object_name)
#                     else:
#                         current_category_id = category_set[object_name]
#
#                 elif current_parent == 'size':
#                     if size[subelem.tag] is not None:
#                         raise Exception('xml structure broken at size tag.')
#                     size[subelem.tag] = int(subelem.text)
#
#                 # option is <xmin>, <ymin>, <xmax>, <ymax>, when subelem is <bndbox>
#                 for option in subelem:
#                     if current_sub == 'bndbox':
#                         if bndbox[option.tag] is not None:
#                             raise Exception('xml structure corrupted at bndbox tag.')
#                         bndbox[option.tag] = int(option.text)
#
#                 # only after parse the <object> tag
#                 if bndbox['xmin'] is not None:
#                     if object_name is None:
#                         raise Exception('xml structure broken at bndbox tag')
#                     if current_image_id is None:
#                         raise Exception('xml structure broken at bndbox tag')
#                     if current_category_id is None:
#                         raise Exception('xml structure broken at bndbox tag')
#                     bbox = []
#                     # x
#                     bbox.append(bndbox['xmin'])
#                     # y
#                     bbox.append(bndbox['ymin'])
#                     # w
#                     bbox.append(bndbox['xmax'] - bndbox['xmin'])
#                     # h
#                     bbox.append(bndbox['ymax'] - bndbox['ymin'])
#                     print('add annotation with {},{},{},{}'.format(object_name, current_image_id, current_category_id,
#                                                                    bbox))
#                     addAnnoItem(object_name, current_image_id, current_category_id, bbox)
#     json.dump(coco, open(json_save_path, 'w'))


"""直接从xml文件夹中生成"""
def parseXmlFiles(xml_path, json_save_path):
    num_img = 0
    num_gt = 0
    for f in os.listdir(xml_path):
        if not f.endswith('.xml'):
            continue

        num_img += 1
        xml_file = os.path.join(xml_path, f)
        print(xml_file)

        tree = ET.parse(xml_file)
        root = tree.getroot()
        if root.tag != 'annotation':
            raise Exception('pascal voc xml root element should be annotation, rather than {}'.format(root.tag))

        ''' 
        先统计图片名字/图片尺寸 
        '''
        file_name = None # 图片名字
        size = dict()
        size['width'] = None
        size['height'] = None
        size['depth'] = None
        for elem in root: # 遍历<annotation>的子标签
            current_parent = elem.tag  # <annotation>子类名

            # 找到文件名字
            if current_parent == 'filename': # 将xml的图片名,存于images字段的file_name中,需包含.jpg
                file_name = elem.text
                file_name = file_name.strip().split(".")[0] + '.jpg' # jpg的图片名
                if file_name in category_set:
                    raise Exception('file_name duplicated')

            # 找到图片尺寸
            if current_parent == 'size':
                for subelem in elem:
                    current_sub = subelem.tag
                    if size[current_sub] is not None: # <depth>,<height>,<width>的值为空才能赋值
                        raise Exception('xml structure broken at size tag.')
                    size[current_sub] = int(subelem.text)

        ''' 
        赋予图片id 
        '''
        current_image_id = None
        # 如果获取了图片名和图片尺寸,则给予该图片对应id
        if file_name is not None and size['width'] is not None and size['height'] is not None \
                        and size['depth'] is not None:
            if file_name not in image_set:  # 图片名不能和之前重复
                current_image_id = addImgItem(file_name, size)  # 赋予图片id
                print('add image with {} and {}'.format(file_name, size))
            else:
                raise Exception('duplicated image: {}'.format(file_name))

        ''' 
        统计每个GT的信息 
        '''
        for elem in root: # 遍历<annotation>的子标签
            current_parent = elem.tag # <annotation>子类名
            object_name = None

            ''' 找到目标 '''
            if current_parent == 'object':
                num_gt += 1

                current_category_id = None # 类别id
                bndbox = dict() # bbox
                bndbox['xmin'] = None
                bndbox['xmax'] = None
                bndbox['ymin'] = None
                bndbox['ymax'] = None
                for subelem in elem: # 遍历<object>子类
                    current_sub = subelem.tag

                    # GT的名字
                    if current_sub == 'name':
                        object_name = subelem.text  # 获取该目标类别名字
                        if object_name not in category_set:  # 定义该类别的id
                            current_category_id = addCatItem(object_name)
                        else:
                            current_category_id = category_set[object_name]

                    # GT的坐标
                    if current_sub == 'bndbox':
                        for option in subelem: # 统计4个坐标<xmin>, <ymin>, <xmax>, <ymax>
                            if bndbox[option.tag] is not None:
                                raise Exception('xml structure corrupted at bndbox tag.')
                            bndbox[option.tag] = int(float(option.text)) # 有的坐标是浮点数,因此需要先float

                if object_name is None: # 该目标得有类别名字
                    raise Exception('goal need class name')
                if current_category_id is None: # 类别id不能为空
                    raise Exception("category id can't be none")
                if current_image_id is None: # 图片id不能空
                    raise Exception("img id can't be none")
                if bndbox['xmin'] is None or bndbox['xmax'] is None or bndbox['ymin'] is None or bndbox['ymax'] is None:
                    raise Exception("bbox can't be none")

                # 得到该目标的json
                bbox = []
                bbox.append(bndbox['xmin'])  # x
                bbox.append(bndbox['ymin'])  # y
                bbox.append(bndbox['xmax'] - bndbox['xmin'])  # w
                bbox.append(bndbox['ymax'] - bndbox['ymin'])  # h
                print('add annotation with {},{},{},{}'.format(object_name, current_image_id, current_category_id, bbox))
                addAnnoItem(object_name, current_image_id, current_category_id, bbox)
    print(f'\nimg num: {num_img}    gt num: {num_gt}')
    json.dump(coco, open(json_save_path, 'w'))


if __name__ == '__main__':
    # Tag: 标签,<width name='yes'>xxx</width>
    # Element:被Tag包围的部分,640或子标签们
    # Attribute：标签的属性,name

    # 仅提供标签即可
    ann_path = "./voc/annotations"
    # ann_path = "./红外飞机voc/annotations"
    json_save_path = "trainval.json"
    parseXmlFiles(ann_path, json_save_path)
